import streamlit as st
import pandas as pd
import json
import os
from pathlib import Path
import sys
import time
from typing import List, Dict, Any, Optional
from src.pipeline import Pipeline, RunConfig

# 添加项目路径到系统路径
project_root = Path(__file__).parent
sys.path.append(str(project_root))
sys.path.append(str(project_root / "src"))

# 导入项目模块
try:
    from src.pipeline import Pipeline, RunConfig
    from src.questions_processing import QuestionsProcessor
    from src.pdf_mineru import set_mode, get_current_mode
except ImportError as e:
    st.error(f"导入模块失败: {e}")
    st.error("请确保在项目根目录运行此应用")
    st.stop()

# 页面配置
st.set_page_config(
    page_title="RAG Challenge 2 - RTX 5080 Powered",
    page_icon="🚀",
    layout="wide",
    initial_sidebar_state="expanded"
)

# 自定义CSS样式
st.markdown("""
<style>
    .main-header {
        background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
        padding: 1rem;
        border-radius: 10px;
        color: white;
        text-align: center;
        margin-bottom: 2rem;
    }
    
    .metric-card {
        background: white;
        padding: 1rem;
        border-radius: 8px;
        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
        border-left: 4px solid #667eea;
    }
    
    .result-card {
        background: #f8f9fa;
        padding: 1rem;
        border-radius: 8px;
        margin: 0.5rem 0;
        border-left: 3px solid #28a745;
    }
    
    .search-box {
        background: white;
        padding: 1.5rem;
        border-radius: 10px;
        box-shadow: 0 2px 8px rgba(0,0,0,0.1);
        margin-bottom: 1rem;
    }
    
    .sidebar-section {
        background: #f8f9fa;
        padding: 1rem;
        border-radius: 8px;
        margin-bottom: 1rem;
    }
</style>
""", unsafe_allow_html=True)

# 初始化会话状态
def init_session_state():
    """初始化会话状态"""
    if 'pipeline' not in st.session_state:
        st.session_state.pipeline = None
    if 'companies' not in st.session_state:
        st.session_state.companies = []
    if 'search_results' not in st.session_state:
        st.session_state.search_results = []
    if 'current_answer' not in st.session_state:
        st.session_state.current_answer = ""
    if 'search_history' not in st.session_state:
        st.session_state.search_history = []

def load_companies() -> List[str]:
    """加载可用的公司列表"""
    try:
        data_path = Path("data/stock_data")
        subset_file = data_path / "subset.csv"
        
        if subset_file.exists():
            # 尝试不同的编码方式读取CSV文件
            encodings = ['utf-8', 'gbk', 'gb18030', 'latin1']
            
            for encoding in encodings:
                try:
                    df = pd.read_csv(subset_file, encoding=encoding)
                    if 'company_name' in df.columns:
                        companies = df['company_name'].unique().tolist()
                        return sorted(companies)
                    else:
                        # 如果没有company_name列，尝试查看所有列
                        st.warning(f"CSV文件中没有'company_name'列，可用列: {df.columns.tolist()}")
                        # 尝试使用第一列作为公司名
                        if len(df.columns) > 0:
                            companies = df[df.columns[0]].unique().tolist()
                            return sorted(companies)
                except UnicodeDecodeError:
                    continue
                except Exception as e:
                    st.warning(f"使用{encoding}编码读取失败: {e}")
                    continue
            
            # 如果所有编码都失败，尝试二进制方式读取并手动解析
            try:
                with open(subset_file, 'rb') as f:
                    content = f.read()
                    # 尝试检测编码
                    import chardet
                    detected = chardet.detect(content)
                    if detected['confidence'] > 0.5:
                        encoding = detected['encoding']
                        st.info(f"检测到文件编码: {encoding}，置信度: {detected['confidence']}")
                        df = pd.read_csv(subset_file, encoding=encoding)
                        if 'company_name' in df.columns:
                            companies = df['company_name'].unique().tolist()
                            return sorted(companies)
            except ImportError:
                st.warning("未安装chardet库，无法自动检测文件编码")
            except Exception as e:
                st.warning(f"尝试自动检测编码失败: {e}")
            
            st.error("无法正确读取CSV文件，请检查文件编码")
            return []
        else:
            st.warning("未找到 subset.csv 文件，请先处理 PDF 文档")
            return []
    except Exception as e:
        st.error(f"加载公司列表失败: {e}")
        return []

def initialize_pipeline(profile_name: str = "default") -> Optional[Pipeline]:
    """初始化处理管道"""
    try:
        root_path = Path("data/stock_data")
        if not root_path.exists():
            st.error("数据目录不存在，请先处理 PDF 文档")
            return None
        
        # 使用新的配置系统创建 Pipeline
        try:
            pipeline = Pipeline.from_config_profile(root_path, profile_name)
            st.success(f"✅ 已加载配置方案: {profile_name}")
            return pipeline
        except Exception as config_error:
            st.warning(f"使用配置方案失败: {config_error}")
            # 回退到默认配置
            pipeline = Pipeline(root_path, run_config=RunConfig())
            st.info("已回退到默认配置")
            return pipeline
        
    except Exception as e:
        st.error(f"初始化管道失败: {e}")
        return None

def search_documents(pipeline: Pipeline, company: str, query: str, config: Dict[str, Any]) -> tuple:
    """搜索文档并生成答案"""
    try:
        # 创建问题处理器 - 修复参数传递
        processor = QuestionsProcessor(
            vector_db_dir=pipeline.paths.vector_db_dir,
            documents_dir=pipeline.paths.documents_dir,
            subset_path=pipeline.paths.subset_path,
            new_challenge_pipeline=True,
            parent_document_retrieval=pipeline.run_config.parent_document_retrieval,
            retrieval_method=pipeline.run_config.retrieval_method,
            llm_reranking=pipeline.run_config.llm_reranking,
            llm_reranking_sample_size=pipeline.run_config.llm_reranking_sample_size,
            top_n_retrieval=config.get('top_k', pipeline.run_config.top_n_retrieval),
            parallel_requests=pipeline.run_config.parallel_requests,
            api_provider=pipeline.run_config.api_provider,
            answering_model=pipeline.run_config.answering_model,
            full_context=pipeline.run_config.full_context
        )
        
        # 执行搜索
        with st.spinner("正在搜索相关文档..."):
            results = processor.process_single_question(
                question=query,
                kind="text"  # 或根据用户选择的答案类型
            )
        
        return results.get('retrieved_documents', []), results.get('answer', '')
    
    except Exception as e:
        st.error(f"搜索失败: {e}")
        return [], ""

def main():
    """主函数"""
    init_session_state()
    
    # 主标题
    st.markdown("""
    <div class="main-header">
        <h1>🚀 RAG-main - RTX 5080 Powered</h1>
        <p>基于 RAG 系统，由 RTX 5080 GPU 加速</p>
        <p>📊 支持多家公司的年报问答 | 🔥 问题检索 + LLM重排序 + GPT-4o</p>
    </div>
    """, unsafe_allow_html=True)
    
    # 侧边栏配置
    with st.sidebar:
        st.markdown('<div class="sidebar-section">', unsafe_allow_html=True)
        st.header("🎯 首页设置")
        
        # 选择公司
        st.subheader("📊 选择公司")
        st.caption("选择要查询的公司年报")
        
        if not st.session_state.companies:
            st.session_state.companies = load_companies()
        
        if st.session_state.companies:
            selected_company = st.selectbox(
                "公司列表",
                st.session_state.companies,
                key="company_select"
            )
        else:
            st.warning("暂无可用公司数据")
            selected_company = None
        
        st.markdown('</div>', unsafe_allow_html=True)
        
        # 问题类型配置
        st.markdown('<div class="sidebar-section">', unsafe_allow_html=True)
        st.subheader("🔧 问题类型")
        st.caption("选择问题的答案类型")
        
        answer_types = ["boolean", "number", "name", "text"]
        selected_type = st.selectbox(
            "答案类型",
            answer_types,
            index=3,  # 默认选择 text
            key="answer_type"
        )
        st.markdown('</div>', unsafe_allow_html=True)
        
        # 高级配置
        st.markdown('<div class="sidebar-section">', unsafe_allow_html=True)
        st.subheader("⚙️ 高级配置")
        
        # 配置方案选择
        try:
            from src.config_manager import get_available_profiles, get_config_manager
            
            available_profiles = get_available_profiles()
            config_manager = get_config_manager()
            
            # 创建配置选项字典，包含描述
            profile_options = {}
            for profile in available_profiles:
                try:
                    description = config_manager.get_profile_description(profile)
                    profile_options[profile] = f"{profile} - {description}"
                except:
                    profile_options[profile] = profile
            
            selected_profile = st.selectbox(
                "🎯 配置方案",
                list(profile_options.keys()),
                format_func=lambda x: profile_options[x],
                key="profile_select",
                help="选择预定义的配置方案，包含模型、检索参数等设置"
            )
            
            # 显示当前配置详情
            if selected_profile:
                try:
                    profile_config = config_manager.get_profile_config(selected_profile)
                    st.info(f"""
                    **当前配置详情:**
                    - 🤖 API 提供商: {profile_config.api_provider}
                    - 🧠 回答模型: {profile_config.answering_model}
                    - 🔄 并行请求: {profile_config.parallel_requests}
                    - 🔍 检索数量: {profile_config.retrieval.top_n_retrieval}
                    - 🎯 LLM 重排序: {'启用' if profile_config.retrieval.llm_reranking else '禁用'}
                    """)
                except Exception as e:
                    st.warning(f"获取配置详情失败: {e}")
            
        except Exception as e:
            st.warning(f"加载配置方案失败: {e}")
            # 回退到原有的配置选择
            config_options = {
                "max": "最大配置 (推荐)",
                "max_docling": "Docling 最大配置", 
                "base": "基础配置",
                "pdr": "父文档检索配置"
            }
            
            selected_profile = st.selectbox(
                "检索配置",
                list(config_options.keys()),
                format_func=lambda x: config_options[x],
                key="config_select"
            )
        
        # MinerU 模式配置
        st.subheader("🔄 MinerU 模式")
        current_mode = get_current_mode()
        st.info(f"当前模式: {current_mode}")
        
        col1, col2 = st.columns(2)
        with col1:
            if st.button("本地模式", type="secondary"):
                set_mode("local")
                st.success("已切换到本地模式")
                st.rerun()
        
        with col2:
            if st.button("云端模式", type="secondary"):
                set_mode("cloud")
                st.success("已切换到云端模式")
                st.rerun()
        
        # 使用 LLM 重排序
        use_reranking = st.checkbox("启用 LLM 重排序", value=True, key="use_reranking")
        
        # 检索参数
        top_k = st.slider("检索文档数量", min_value=3, max_value=20, value=6, key="top_k")
        
        st.markdown('</div>', unsafe_allow_html=True)
    
    # 主内容区域
    col1, col2 = st.columns([1, 1])
    
    with col1:
        st.markdown('<div class="search-box">', unsafe_allow_html=True)
        st.subheader("❓ 输入问题")
        st.caption("输入您想要查询的问题")
        
        # 预设问题示例
        example_questions = [
            "Who is the CEO of the company?",
            "公司的主要业务是什么？",
            "公司的营收情况如何？",
            "公司面临的主要风险有哪些？",
            "公司的研发投入占比是多少？"
        ]
        
        selected_example = st.selectbox(
            "选择示例问题（可选）",
            [""] + example_questions,
            key="example_question"
        )
        
        # 问题输入框
        if selected_example:
            default_question = selected_example
        else:
            default_question = ""
        
        user_question = st.text_area(
            "您的问题",
            value=default_question,
            height=100,
            placeholder="请输入您的问题...",
            key="user_question"
        )
        
        # 搜索按钮
        search_button = st.button("🔍 搜索文档", type="primary", use_container_width=True)
        
        st.markdown('</div>', unsafe_allow_html=True)
        
        # 搜索历史
        if st.session_state.search_history:
            st.subheader("📚 搜索历史")
            for i, (q, company, timestamp) in enumerate(reversed(st.session_state.search_history[-5:])):
                with st.expander(f"{timestamp} - {company}"):
                    st.write(q)
                    if st.button(f"重新搜索", key=f"research_{i}"):
                        st.session_state.user_question = q
                        st.session_state.company_select = company
                        st.rerun()
    
    with col2:
        st.subheader("📊 检索结果")
        
        # 执行搜索
        if search_button and user_question and selected_company:
            # 初始化管道
            if st.session_state.pipeline is None or st.session_state.get('current_profile') != selected_profile:
                st.session_state.pipeline = initialize_pipeline(selected_profile)
                st.session_state.current_profile = selected_profile
            
            if st.session_state.pipeline:
                # 执行搜索
                documents, answer = search_documents(
                    st.session_state.pipeline,
                    selected_company,
                    user_question,
                    {"top_k": top_k, "use_reranking": use_reranking}
                )
                
                st.session_state.search_results = documents
                st.session_state.current_answer = answer
                
                # 添加到搜索历史
                timestamp = time.strftime("%H:%M:%S")
                st.session_state.search_history.append((user_question, selected_company, timestamp))
        
        # 显示答案
        if st.session_state.current_answer:
            st.markdown("### 🎯 生成答案")
            st.markdown(f'<div class="result-card">{st.session_state.current_answer}</div>', 
                       unsafe_allow_html=True)
        
        # 显示检索结果
        if st.session_state.search_results:
            st.markdown("### 📄 相关文档片段")
            
            for i, doc in enumerate(st.session_state.search_results[:5]):  # 显示前5个结果
                with st.expander(f"结果 {i+1} (相似度: {doc.get('score', 0):.3f})"):
                    # 显示文档信息
                    col_a, col_b = st.columns([1, 1])
                    with col_a:
                        st.write(f"**页码**: {doc.get('page', 'N/A')}")
                        st.write(f"**公司**: {doc.get('company_name', 'N/A')}")
                    with col_b:
                        st.write(f"**文件**: {doc.get('file_name', 'N/A')}")
                        st.write(f"**行数**: {doc.get('lines', 'N/A')}")
                    
                    # 显示文档内容
                    st.markdown("**内容**:")
                    content = doc.get('text', '').strip()
                    if len(content) > 500:
                        st.write(content[:500] + "...")
                        if st.button(f"显示完整内容", key=f"show_full_{i}"):
                            st.write(content)
                    else:
                        st.write(content)
    
    # 底部信息
    st.markdown("---")
    col1, col2, col3 = st.columns(3)
    
    with col1:
        if st.session_state.companies:
            st.metric("可用公司数量", len(st.session_state.companies))
    
    with col2:
        if st.session_state.search_results:
            st.metric("检索结果数量", len(st.session_state.search_results))
    
    with col3:
        st.metric("搜索历史", len(st.session_state.search_history))

if __name__ == "__main__":
    main()